from torch.utils.data import Dataset
import os
from skimage import io
import torch
from torchvision import transforms

labels = []
#2-9
for i in range(8):
    labels.append(50 + i)
for i in range(26):
    labels.append(65 + i)

class VerCodeDataset(Dataset):
    def __init__(self, image_dir="./letter_template/"):
        l = os.listdir(image_dir)
        self.data = []
        self.label = []
        for d in l:
            fs = os.listdir("{}{}".format(image_dir, d))
            for f in fs:
                fup = "{}{}/{}".format(image_dir, d, f)
                t = torch.from_numpy(io.imread(fup)).float() / 255
                norl = transforms.Normalize(t.mean(), t.std())
                self.data.append(norl(t.reshape(1, 40, 40)))
                #self.label.append(d)
                self.label.append(labels.index(ord(d)))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
        return {"data": self.data[item], "label": self.label[item]}

from torch.utils.data import DataLoader

def get_labels():
    return labels

def trainloader(bs):
    ds = VerCodeDataset()
    return DataLoader(ds,shuffle=True,batch_size=bs)

def testloader():
    ds = VerCodeDataset(image_dir="./letter_test/")
    print(labels)
    return DataLoader(ds,batch_size=5)

if (__name__ == "__main__"):
    tl = trainloader(5)
    for step,i in enumerate(tl):
        print(step)
        print(i["data"])
        print(i["label"])
        exit(0)